#include "irbuilder.h"

#include <functional>
#include <iostream>
#include <ostream>
#include <string>
#include <unordered_set>
#include <vector>

#include "analyzer.h"
#include "arithmetic.h"
#include "ast.h"
#include "irfunction.h"
#include "iroperand.h"
#include "semantics.h"
#include "symbol.h"
#include "token.h"
#include "types.h"
#include "utils.h"

#define DECL_VISIT(className) \
    ASTNode* IRBuilder::Visit##className(className* node)

using namespace Utils;
using namespace AST;
using namespace Semantics;

IR::Variable* IRBuilder::VarMap::Get(Symbol* sym) {
    if (mp.count(sym)) return new IR::Variable(*mp[sym]);
    bool isPtr = sym->info->As<Semantics::Pointer>();
    auto var = new IR::Variable(sym->info->baseType, isPtr);
    var->isGlobal = sym->isGlobal;
    var->name = std::to_string(++count) + "." + sym->name;
    if (isPtr) {
        var->memBase = var->Copy()->As<IR::Variable>();
        var->memBase->memBase = var->memBase;
    }
    return mp[sym] = var;
}

IR::Variable* IRBuilder::VarMap::Get(Type::Basic type, bool isPtr) {
    if (type == Type::Basic::Void) return nullptr;
    auto var = new IR::Variable(type, isPtr);
    var->isGlobal = false;
    var->name = std::to_string(++count);
    return var;
}

void CondConnect(ASTNode* node) {
    if (node->irTPos && node->irFPos) {
        auto zero = new IR::Immediate(node->irValue->type.base);
        node->irEnd->Connect(ArithOp::NotEqual, node->irValue, zero,
                             node->irTPos, node->irFPos);
    }
}

DECL_VISIT(AssignExpr) {
    if (auto addr = node->lhs->As<AddrExpr>()) {
        Visit(addr->base);
        Visit(addr->offset);
        Visit(node->rhs);

        auto store = new IR::Store();
        store->base = addr->base->irValue->Copy()->As<IR::Variable>();
        store->offset = addr->offset->irValue->Copy();
        store->src = node->rhs->irValue->Copy();
        node->irEnd = new IRBlock(store);

        node->irBegin = addr->base->irBegin;
        addr->base->irEnd->Connect(addr->offset->irBegin);
        addr->offset->irEnd->Connect(node->rhs->irBegin);
        node->rhs->irEnd->Connect(node->irEnd);
        node->irValue = node->rhs->irValue;
        return node;
    }

    Visit(node->lhs);
    Visit(node->rhs);

    auto assi = new IR::Operation();
    assi->dest = node->lhs->irValue->Copy()->As<IR::Variable>();
    assi->rhs = node->rhs->irValue->Copy();
    assi->op = ArithOp::Add;
    node->irEnd = new IRBlock(assi);

    node->irBegin = node->lhs->irBegin;
    node->lhs->irEnd->Connect(node->rhs->irBegin);
    node->rhs->irEnd->Connect(node->irEnd);
    node->irValue = node->rhs->irValue;

    CondConnect(node);
    return node;
}

DECL_VISIT(CallExpr) {
    auto call = new IR::Call();
    auto funcInfo = node->func->As<IdentifierExpr>()->info;
    Assert(funcMap.count(funcInfo), "unrecognized function");
    call->func = funcMap[funcInfo];
    call->dest = varMap.Get(funcInfo->baseType);

    node->irBegin = node->irEnd = new IRBlock(call);
    auto args = node->args->As<ArgList>()->items;
    std::reverse(args.begin(), args.end());
    for (auto arg : args) {
        Visit(arg);
        arg->irEnd->Connect(node->irBegin);
        node->irBegin = arg->irBegin;
        call->args.push_back(arg->irValue->Copy());
    }
    std::reverse(call->args.begin(), call->args.end());
    node->irValue = call->dest;

    CondConnect(node);
    return node;
}

DECL_VISIT(ArrayExpr) {
    Err("");
    return node;
}

DECL_VISIT(OperatorExpr) {
    if (node->op == ArithOp::LogicalAnd) {
        node->rhs->irTPos = node->irTPos;
        node->rhs->irFPos = node->irFPos;
        Visit(node->rhs);
        node->lhs->irTPos = node->rhs->irBegin;
        node->lhs->irFPos = node->irFPos;
        Visit(node->lhs);
        node->irBegin = node->lhs->irBegin;
        node->irEnd = node->rhs->irEnd;
        return node;
    }
    if (node->op == ArithOp::LogicalOr) {
        node->rhs->irTPos = node->irTPos;
        node->rhs->irFPos = node->irFPos;
        Visit(node->rhs);
        node->lhs->irTPos = node->irTPos;
        node->lhs->irFPos = node->rhs->irBegin;
        Visit(node->lhs);
        node->irBegin = node->lhs->irBegin;
        node->irEnd = node->rhs->irEnd;
        return node;
    }
    if (node->op == ArithOp::Not) {
        node->rhs->irTPos = node->irFPos;
        node->rhs->irFPos = node->irTPos;
        Visit(node->rhs);
        node->irBegin = node->rhs->irBegin;
        node->irEnd = node->rhs->irEnd;
        return node;
    }

    Visit(node->lhs);
    Visit(node->rhs);

    if (node->op == ArithOp::Less || node->op == ArithOp::LessEqual ||
        node->op == ArithOp::Greater || node->op == ArithOp::GreaterEqual ||
        node->op == ArithOp::Equal || node->op == ArithOp::NotEqual) {
        auto lhs = node->lhs->irValue->Copy();
        auto rhs = node->rhs->irValue->Copy();

        node->irBegin = node->lhs->irBegin;
        node->lhs->irEnd->Connect(node->rhs->irBegin);
        node->irEnd = node->rhs->irEnd;
        Assert(!node->irEnd->GetTerminator(), "");
        node->irEnd->Connect(node->op, lhs, rhs, node->irTPos, node->irFPos);
        return node;
    }

    auto ir = new IR::Operation();
    if (node->lhs->irValue) ir->lhs = node->lhs->irValue->Copy();
    ir->rhs = node->rhs->irValue->Copy();
    ir->op = node->op;
    ir->dest = varMap.Get(node->info->baseType);
    node->irEnd = new IRBlock(ir);

    node->irBegin = node->lhs->irBegin;
    node->lhs->irEnd->Connect(node->rhs->irBegin);
    node->rhs->irEnd->Connect(node->irEnd);
    node->irValue = ir->dest;

    CondConnect(node);
    return node;
}

DECL_VISIT(TernaryExpr) {
    Visit(node->tVal);
    Visit(node->fVal);

    node->irEnd = new IRBlock();
    node->irValue = varMap.Get(node->info->baseType);

    auto tAssi = new IR::Operation();
    tAssi->dest = node->irValue->Copy()->As<IR::Variable>();
    tAssi->rhs = node->tVal->irValue->Copy();
    tAssi->op = ArithOp::Add;
    auto tBlk = new IRBlock(tAssi);
    node->tVal->irEnd->Connect(tBlk);
    tBlk->Connect(node->irEnd);

    auto fAssi = new IR::Operation();
    fAssi->dest = node->irValue->Copy()->As<IR::Variable>();
    fAssi->rhs = node->fVal->irValue->Copy();
    fAssi->op = ArithOp::Add;
    auto fBlk = new IRBlock(fAssi);
    node->fVal->irEnd->Connect(fBlk);
    fBlk->Connect(node->irEnd);

    node->cond->irTPos = node->tVal->irBegin;
    node->cond->irFPos = node->fVal->irBegin;
    Visit(node->cond);
    node->irBegin = node->cond->irBegin;

    CondConnect(node);
    return node;
}

DECL_VISIT(LiteralExpr) {
    Assert(node->info->As<Semantics::Constant>(), "");
    node->irBegin = node->irEnd = new IRBlock();
    auto* info = node->info->As<Semantics::Constant>();
    node->irValue = new IR::Immediate(info->value);

    CondConnect(node);
    return node;
}

DECL_VISIT(VoidExpr) {
    node->irBegin = node->irEnd = new IRBlock();
    return node;
}

DECL_VISIT(IdentifierExpr) {
    node->irBegin = node->irEnd = new IRBlock();
    IR::Variable* var;
    if (auto info = node->info->As<Semantics::LValue>()) {
        var = varMap.Get(info->symbol);
    } else if (auto info = node->info->As<Semantics::Pointer>()) {
        var = varMap.Get(info->symbol);
    } else
        Err("");
    node->irValue = var;

    CondConnect(node);
    return node;
}

DECL_VISIT(InitializerExpr) {
    Err("");
    return node;
}

DECL_VISIT(CastExpr) {
    Visit(node->origin);

    auto ir = new IR::Operation();
    ir->rhs = node->origin->irValue->Copy();
    bool toInt = node->to == Type::Basic::Int;
    ir->op = toInt ? ArithOp::IntCast : ArithOp::FloatCast;
    ir->dest = varMap.Get(node->to);
    node->irEnd = new IRBlock(ir);

    node->irBegin = node->origin->irBegin;
    node->origin->irEnd->Connect(node->irEnd);
    node->irValue = ir->dest;

    CondConnect(node);
    return node;
}

DECL_VISIT(AddrExpr) {
    Visit(node->base);
    Visit(node->offset);

    if (node->info->As<Pointer>()) {
        auto add = new IR::Operation();
        add->lhs = node->base->irValue->Copy();
        add->rhs = node->offset->irValue->Copy();
        add->op = ArithOp::Add;
        add->dest = varMap.Get(node->info->baseType, true);
        add->dest->memBase = add->lhs->Copy()->As<IR::Variable>();
        // std::cout << add->dest->memBase << std::endl;
        node->irEnd = new IRBlock(add);
        node->irValue = add->dest;
    } else {
        auto load = new IR::Load();
        load->base = node->base->irValue->Copy()->As<IR::Variable>();
        load->offset = node->offset->irValue->Copy();
        load->dest = varMap.Get(node->info->baseType);
        node->irEnd = new IRBlock(load);
        node->irValue = load->dest;
    }

    node->irBegin = node->base->irBegin;
    node->base->irEnd->Connect(node->offset->irBegin);
    node->offset->irEnd->Connect(node->irEnd);

    CondConnect(node);
    return node;
}

DECL_VISIT(ArgList) {
    Err("");
    return node;
}

DECL_VISIT(Block) {
    auto defs = new IRBlock();
    for (auto sym : node->local) {
        if (!sym->info->As<Semantics::Pointer>()) continue;
        auto alloc = new IR::Alloc();
        alloc->dest = varMap.Get(sym);
        alloc->size = sym->allocMem;
        defs->irList.emplace_back(alloc);
    }
    node->irBegin = node->irEnd = defs;
    for (auto item : node->items) {
        Visit(item);
        node->irEnd->Connect(item->irBegin);
        node->irEnd = item->irEnd;
    }
    return node;
}

DECL_VISIT(VarDecl) {
    Err("");
    return node;
}

DECL_VISIT(FuncDecl) {
    currentFunc->name = node->name;
    for (auto arg : node->args) {
        Visit(arg);
        currentFunc->args.push_back(arg->irValue->Copy()->As<IR::Variable>());
    }
    currentFunc->ret = varMap.Get(node->retType);
    auto irRet = new IR::Return();
    irRet->ret = currentFunc->ret;
    currentFunc->exit = new IRBlock(irRet);

    Visit(node->content);
    currentFunc->entry = node->content->irBegin;
    node->content->irEnd->Connect(currentFunc->exit);
    return node;
}

DECL_VISIT(IfStmt) {
    Visit(node->tVal);
    Visit(node->fVal);

    node->irEnd = new IRBlock();
    node->tVal->irEnd->Connect(node->irEnd);
    node->fVal->irEnd->Connect(node->irEnd);

    node->cond->irTPos = node->tVal->irBegin;
    node->cond->irFPos = node->fVal->irBegin;
    Visit(node->cond);
    node->irBegin = node->cond->irBegin;
    return node;
}

DECL_VISIT(WhileStmt) {
    // do-while section
    node->cond->irTPos = new IRBlock();
    node->cond->irFPos = node->irEnd = new IRBlock();
    Visit(node->cond);

    irBreak.push_back(node->irEnd);
    irContinue.push_back(node->cond->irBegin);

    Visit(node->content);
    node->cond->irTPos->Connect(node->content->irBegin);
    node->content->irEnd->Connect(node->cond->irBegin);

    // if section
    node->cond->irTPos = node->content->irBegin;
    node->cond->irFPos = node->irEnd;
    Visit(node->cond);
    node->irBegin = node->cond->irBegin;

    irBreak.pop_back();
    irContinue.pop_back();
    return node;
}

DECL_VISIT(ContinueStmt) {
    Assert(!irContinue.empty(), "unexpected continue");
    auto block = new IRBlock();
    block->Connect(irContinue.back());
    node->irBegin = node->irEnd = block;
    return node;
}

DECL_VISIT(BreakStmt) {
    Assert(!irBreak.empty(), "unexpected break");
    auto block = new IRBlock();
    block->Connect(irBreak.back());
    node->irBegin = node->irEnd = block;
    return node;
}

DECL_VISIT(ReturnStmt) {
    Visit(node->ret);

    if (currentFunc->ret) {
        auto assi = new IR::Operation();
        assi->dest = currentFunc->ret->Copy()->As<IR::Variable>();
        assi->rhs = node->ret->irValue->Copy();
        assi->op = ArithOp::Add;
        node->irEnd = new IRBlock(assi);

        node->irBegin = node->ret->irBegin;
        node->ret->irEnd->Connect(node->irEnd);
        node->irEnd->Connect(currentFunc->exit);
    } else {
        node->irBegin = node->irEnd = new IRBlock();
        node->irEnd->Connect(currentFunc->exit);
    }
    return node;
}

void IRBuilder::Visit(ASTNode*& node) { ASTVisitor::Visit(node); }

void* IRBuilder::Run(void* input) {
    auto symbolMap = (SymbolMap*)input;
    auto module = new IRModule();
    for (auto [name, sym] : symbolMap->ext) {
        auto func = sym->info->As<Semantics::Function>();
        Assert(func, "");
        funcMap[func] = currentFunc = new IRFunction();
        currentFunc->name = name;
    }
    for (auto& name : symbolMap->order) {
        auto sym = symbolMap->global[name];
        auto info = sym->info;
        if (auto func = info->As<Semantics::Function>()) {
            funcMap[func] = currentFunc = new IRFunction();
            Visit(func->func);
            currentFunc->Simplify();
            currentFunc->SetParent();
            currentFunc->Validate();
            module->funcs[name] = currentFunc;
            module->globalFuncs.push_back(currentFunc);
        } else {
            module->globalVars[varMap.Get(sym)] = sym;
        }
    }
    return module;
}
